#include "dominator_tree.h"
#include "../../include/ir.h"
#include <vector>
#include <iostream>
void DominatorTree::PrintBlock(const LLVMBlock& block) const {
    std::cout << "Block " << block->block_id;
}
void DominatorTree::PrintDomTree() const {
    std::cout << "Dom Tree:\n";
    for (size_t i = 0; i < dom_tree.size(); ++i) {
        std::cout << "Node " << i << ": ";
        for (const auto& block : dom_tree[i]) {
            PrintBlock(block);
            std::cout << " ";
        }
        std::cout << "\n";
    }
}
void DominatorTree::PrintIdom() const {
    std::cout << "Immediate Dominators (idom):\n";
    for (size_t i = 0; i < idom.size(); ++i) {
        if (i < C->max_label + 1) { // 确保索引有效
            std::cout << "Node " << i << " is immediately dominated by Block ";
            PrintBlock(idom[i]);
            std::cout << "\n";
        }
    }
}
void DominatorTree::PrintDF() const {
    std::cout << "Dominance Frontiers (df):\n";
    for (size_t i = 0; i < df.size(); ++i) {
        std::cout << "Node " << i << ": { ";
        for (int j : df[i]) {
            std::cout << j << " ";
        }
        std::cout << "}\n";
    }
}
void DominatorTree::printdom() const {
    if (dom.empty()) {
        std::cout << "Dom matrix is empty." << std::endl;
        return;
    }

    size_t num_nodes = dom.size();
    std::cout << "Dom Matrix (dom):\n";
    std::cout << "  ";
    for (size_t i = 0; i < num_nodes; ++i) {
        std::cout << i << " ";
    }
    std::cout << "\n";

    for (size_t i = 0; i < num_nodes; ++i) {
        std::cout << i << " ";
        for (size_t j = 0; j < num_nodes; ++j) {
            if (j < dom[i].size()) {
                std::cout << (dom[i][j] ? "1" : "0") << " ";
            } else {
                std::cout << "? ";
            }
        }
        std::cout << "\n";
    }
}
void DominatorTree::PrintDom() const {
    std::cout << "Immediate Dominators List (Dom):\n";
    for (size_t i = 0; i < Dom.size(); ++i) {
        std::cout << "Node " << i << ": { ";
        for (int j : Dom[i]) {
            std::cout << j << " ";
        }
        std::cout << "}\n";
    }
}

void DominatorTree::PrintDomMatrix() const {
    std::cout << "Dominance Matrix (dom):\n";
    for (size_t i = 0; i < dom.size(); ++i) {
        for (size_t j = 0; j < dom[i].size(); ++j) {
            std::cout << dom[i][j] << " ";
        }
        std::cout << "\n";
    }
}

void DominatorTree::PrintPredecessors() const {
    std::cout << "Predecessors (pre):\n";
    for (size_t i = 0; i < pre.size(); ++i) {
        std::cout << "Node " << i << ": { ";
        for (int j : pre[i]) {
            std::cout << j << " ";
        }
        std::cout << "}\n";
    }
}

void DominatorTree::PrintOrder() const {
    std::cout << "Reverse Post Order (ord):\n";
    for (int i : ord) {
        std::cout << i << " ";
    }
    std::cout << "\n";
}

void DominatorTree::PrintIdomInt() const {
    std::cout << "Immediate Dominators (idom_int):\n";
    for (size_t i = 0; i < idom_int.size(); ++i) {
        std::cout << "Node " << i << " is immediately dominated by Node " << idom_int[i] << "\n";
    }
}

void DominatorTree::PrintEdges() const {
    std::cout << "Edges (e):\n";
    for (size_t i = 0; i < e.size(); ++i) {
        std::cout << "From Node " << i << ": { ";
        for (int j : e[i]) {
            std::cout << j << " ";
        }
        std::cout << "}\n";
    }
}
// 深度优先搜索以获取逆后序排列
//reference https://github.com/optimize\cfg\dominator_tree.cc  10-19
void dfs_postorder(int cur, const std::vector<std::vector<LLVMBlock>> &G, std::vector<int> &result,
                   std::vector<int> &vsd) {
    vsd[cur] = 1; // 标记当前节点已访问
    for (auto next_block : G[cur]) { // 遍历相邻节点
        if (vsd[next_block->block_id] == 0) { // 如果相邻节点未被访问，则递归调用
            dfs_postorder(next_block->block_id, G, result, vsd);
        }
    }
    result.push_back(cur); // 将当前节点加入结果中
}
void DominatorTree::build_dom_tree() {
    int num_nodes = dom.size();
    
    // 确保 dom_tree 的大小与节点数一致，并初始化为空
    dom_tree.resize(num_nodes);
    idom.clear();
    idom.resize(C->max_label+1);
    // 遍历 Dom 列表，构建支配树
    for (int u = 0; u < num_nodes; ++u) {
        for (int v : Dom[u]) {
            if (C->block_map && C->block_map->find(u) != C->block_map->end() &&
                C->block_map->find(v) != C->block_map->end()) {
                // 将 u 添加到 v 的子节点列表中
                idom[u] = (*C->block_map)[v];
                dom_tree[v].push_back((*C->block_map)[u]);
            } else {
                // 处理无效的 block_id 或 block_map 为空的情况
                std::cerr << "Warning: Invalid block_id or block_map is empty." << std::endl;
            }
        }
    }
    /*
    // 特殊处理根节点（假设根节点的 ID 为 0）
    if (!Dom[0].empty()) {
        int root_idom = Dom[0][0];
        if (C->block_map && C->block_map->find(root_idom) != C->block_map->end()) {
            // 根节点的直接支配节点是自己
            dom_tree[root_idom].push_back((*C->block_map)[0]);
        }
    }*/
}
void DominatorTree::getdom() {
    int N=C->max_label+1;   
    bool flag = true; // 标志变量，用于指示是否需要继续迭代
    while (flag) { // 固定点迭代，直到支配关系不再变化
        flag = false;
        
        for (int u : ord) { // 遍历逆后序排列中的每一个节点

            std::vector<bool> tmp(N, false);
            
            // 对于节点 u 的每一个前驱节点 v，计算交集以确定支配 u 的节点;一个点的支配点的点集为它所有前驱结点的支配点集的交集，再并上它本身。
            if (!pre[u].empty()) { // 确保有前驱节点
                tmp = dom[pre[u][0]]; // 初始化为第一个前驱节点的支配集合
                for (size_t i =1; i < pre[u].size(); ++i) {
                    int v = pre[u][i];
                    for (int j = 0; j < N; ++j) {
                        tmp[j] = tmp[j] && dom[v][j]; // 取所有前驱节点支配集合的交集
                    }
                }
            }
            tmp[u] = true; // 确保节点 u 支配自身
            // 如果新的支配集合与旧的不同，则更新并设置标志
            if (tmp != dom[u]) {
                dom[u] = tmp;
                flag = true; // 设置标志表示有更新，需继续迭代
            }
        }
    }

    // 将支配关系转换为立即支配关系列表
    for (int i = 1; i < N; ++i) { // 遍历除起始节点外的所有节点
        Dom[i].clear(); // 清空当前节点的立即支配节点列表
        for (int j = 0; j < N; ++j) { // 检查每一个可能的支配节点
        if(i==j) continue;
            if (dom[i][j]) { // 如果 j 支配 i
                // 找到立即支配节点
                bool isImmediate = true;
                for (int k = 0; k < N; ++k) {
                    if (k != j &&k!=i&& dom[i][k] && dom[k][j]) {// 如果j也支配k且k支配i
                        isImmediate = false;
                        break;
                    }
                }
                if (isImmediate) {
                    Dom[i].push_back(j); // 添加到 i 的立即支配节点列表中
                }
            }
        }
    }
}

// 示例初始化函数
void DominatorTree::initialize() {
    //std::cout<<"initialize"<<std::endl;
    int N=C->max_label+1;
    df.clear();
    df.assign(N, std::vector<int>());
    dom.clear();
    dom.assign(N, std::vector<bool>(N, true)); // 初始化支配关系矩阵
    Dom.clear();
    Dom.assign(N, std::vector<int>()); // 初始化 Dom 数组
    idom_int.clear();
    idom_int.assign(N, -1);
    //初始化ord
    std::vector<int> vsd(N, 0);
    for (int i = 0; i <= C->max_label; i++) {
        vsd.push_back(0);
    }
    dfs_postorder(0,C->G,ord,vsd);
    std::reverse(ord.begin(), ord.end());
    // 初始化 pre 数组
    pre.assign(N, std::vector<int>()); 
    // 遍历 CFG 中的所有基本块
    for (const auto& pair : *C->block_map) {
        int bbid = pair.first;
        LLVMBlock B = pair.second;

        // 获取该基本块的所有前驱节点
        std::vector<LLVMBlock> predecessors = C->GetPredecessor(B);
        for (const auto& pred : predecessors) {
            pre[bbid].push_back(pred->block_id);
        }
    }
    
}

void DomAnalysis::Execute() {
    for (auto [defI, cfg] : llvmIR->llvm_cfg) {
        DomInfo[cfg].C=cfg;
        DomInfo[cfg].BuildDominatorTree();
    }
}
//没有用
void DominatorTree::getidom() {
    int N = C->max_label + 1;
    idom.resize(N, LLVMBlock()); // 初始化 idom 向量，大小为 N
    for (int i = 0; i < Dom.size(); ++i) {
        for (int j = 0; j < Dom[i].size(); ++j) {
            idom[i]=(*C->block_map)[Dom[i][j]];
        }
    }
}

//也没用
void DominatorTree::BuildDomTree() {
        // 确保 N 是正确的
        int N = C->max_label + 1;

        // 初始化 dom_tree
        dom_tree.assign(N, std::vector<LLVMBlock>());

        // 遍历所有基本块并构建支配树
        for (int u = 0; u < N; ++u) {
            if (idom_int[u] != -1) { // 如果有直接支配节点
                LLVMBlock parent_block = (*C->block_map)[idom_int[u]];
                LLVMBlock child_block = (*C->block_map)[u];
                dom_tree[idom_int[u]].push_back(child_block);
            }
        }

        // 处理根节点（假设根节点是 0）
        if (C->block_map->count(0)) {
            LLVMBlock root_block = (*C->block_map)[0];
            dom_tree[0].push_back(root_block); // 根节点自己指向自己
        }
    }
void DominatorTree::initialize_idom() {
    int num_nodes = Dom.size();
    
    // 确保 idom 的大小与 Dom 的大小一致，并初始化为 nullptr
    idom.resize(num_nodes, nullptr);

    // 遍历 Dom 列表，填充 idom
    for (int u = 0; u < num_nodes; ++u) {
        if (!Dom[u].empty()) {
            int v = Dom[u][0]; // 假设每个节点只有一个直接支配节点
            if (C->block_map && C->block_map->find(v) != C->block_map->end()) {
                idom[u] = (*C->block_map)[v];
            } else {
                // 处理无效的 block_id 或 block_map 为空的情况
                std::cerr << "Warning: Invalid block_id " << v << " or block_map is empty." << std::endl;
                idom[u] = nullptr; // 或者采取其他措施
            }
        } else {
            // 如果没有直接支配节点（如根节点），可以设置为 nullptr 或特殊处理
            idom[u] = nullptr;
        }
    }
}
void DominatorTree::BuildDominatorTree(bool reverse) {
    assert(C && "CFG pointer is null");
    auto const *G = &(C->G); // 获取正向图引用
    auto const *invG = &(C->invG); // 获取反向图引用
    auto begin_id = 0; // 开始节点 ID，默认为 0
    if (reverse) { // 如果要构建后支配树
        auto temp = G;
        G = invG;
        invG = temp; // 交换正向图与反向图引用
        Assert(C->ret_block != nullptr); // 断言返回块不为空
        begin_id = C->ret_block->block_id; // 设置开始节点为返回块 ID
    }
   //int block_num = C->max_label + 1; // 计算总共有多少个基本块
    initialize();
    getdom();//求支配关系
    PrintDom();
    //getidom();
    //PrintIdomInt();//不要用，废的
    //initialize_idom();
    build_dom_tree();//建树
    PrintDomTree();
    //BuildDomTree();//不要用，废的
    //idom.resize(block_num); // 初始化立即支配节点列表  
    //reference https://github.com/optimize\cfg\dominator_tree.cc  124-144
    for (int i = 0; i < (*G).size(); i++) {
        for (auto edg_end : (*G)[i]) { // 遍历每条边
            int a = i;
            int b = edg_end->block_id;
            int x = a;
            //这里x=i不就行了，a 的意义何在啊，但是不敢改，哎
            while (x == b || !IsDominate(x, b)) { // 当 x 是 b 或者 x 不支配 b
                df[x].push_back(b); // 设置支配边界
                if (idom[x] != NULL) {
                    //x是b的前驱
                    x = idom[x]->block_id; // 移动到 x 的立即支配节点
                } else {
                    break;
                }
            }
        }
    }
    PrintDF();
}


std::set<int> DominatorTree::GetDF(std::set<int> S) {
    std::set<int> result;
    int N=C->max_label+1;
    for (int id : S) {
        if (id >= 0 && id < N) {
            result.insert(df[id].begin(), df[id].end());
        }
    }

    return result;
}
//计算指定节点集合的支配边界
//支配边界是指在控制流图中，某个基本块的支配树中包含的所有基本块，这些基本块在控制流图中是该基本块的后继，但并不直接被该基本块支配。
std::set<int> DominatorTree::GetDF(int id) {
    std::set<int> result;
    int N=C->max_label+1;
    if (id >= 0 && id < N) {
        result.insert(df[id].begin(), df[id].end());
    }
    return result;
}
//两个块是不是具有支配关系
bool DominatorTree::IsDominate(int id1, int id2) {
    return dom[id2][id1];
}
